{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "S3_PATH_ENTITIES = {\n",
    "    'bert': {\n",
    "        'model': 'v34/entity/bert-base-entity.pb',\n",
    "        'vocab': 'tokenizer/sp10m.cased.bert.vocab',\n",
    "        'tokenizer': 'tokenizer/sp10m.cased.bert.model',\n",
    "        'setting': 'bert-bahasa/dictionary-entities.json',\n",
    "    },\n",
    "    'tiny-bert': {\n",
    "        'model': 'v34/entity/tiny-bert-entity.pb',\n",
    "        'vocab': 'tokenizer/sp10m.cased.bert.vocab',\n",
    "        'tokenizer': 'tokenizer/sp10m.cased.bert.model',\n",
    "        'setting': 'bert-bahasa/dictionary-entities.json',\n",
    "    },\n",
    "    'albert': {\n",
    "        'model': 'v34/entity/albert-base-entity.pb',\n",
    "        'vocab': 'tokenizer/sp10m.cased.v10.vocab',\n",
    "        'tokenizer': 'tokenizer/sp10m.cased.v10.model',\n",
    "        'setting': 'bert-bahasa/dictionary-entities.json',\n",
    "    },\n",
    "    'tiny-albert': {\n",
    "        'model': 'v34/entity/albert-tiny-entity.pb',\n",
    "        'vocab': 'tokenizer/sp10m.cased.v10.vocab',\n",
    "        'tokenizer': 'tokenizer/sp10m.cased.v10.model',\n",
    "        'setting': 'bert-bahasa/dictionary-entities.json',\n",
    "    },\n",
    "    'xlnet': {\n",
    "        'model': 'v34/entity/xlnet-base-entity.pb',\n",
    "        'vocab': 'tokenizer/sp10m.cased.v9.vocab',\n",
    "        'tokenizer': 'tokenizer/sp10m.cased.v9.model',\n",
    "        'setting': 'bert-bahasa/dictionary-entities.json',\n",
    "    },\n",
    "    'alxlnet': {\n",
    "        'model': 'v34/entity/alxlnet-base-entity.pb',\n",
    "        'vocab': 'tokenizer/sp10m.cased.v9.vocab',\n",
    "        'tokenizer': 'tokenizer/sp10m.cased.v9.model',\n",
    "        'setting': 'bert-bahasa/dictionary-entities.json',\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bert\n",
      "tiny-bert\n",
      "albert\n",
      "tiny-albert\n",
      "xlnet\n",
      "alxlnet\n"
     ]
    }
   ],
   "source": [
    "for k in S3_PATH_ENTITIES.keys():\n",
    "    if k != 'multinomial':\n",
    "        print(k)\n",
    "        os.system(f\"wget https://f000.backblazeb2.com/file/malaya-model/{S3_PATH_ENTITIES[k]['model']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.tools.graph_transforms import TransformGraph\n",
    "from glob import glob\n",
    "tf.set_random_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['xlnet-base-entity.pb',\n",
       " 'alxlnet-base-entity.pb',\n",
       " 'albert-tiny-entity.pb',\n",
       " 'tiny-bert-entity.pb',\n",
       " 'bert-base-entity.pb',\n",
       " 'albert-base-entity.pb']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pbs = glob('*.pb')\n",
    "pbs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-8-648ce936170b>:11: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.gfile.GFile.\n",
      "xlnet-base-entity.pb ['Placeholder', 'Placeholder_1', 'Placeholder_2']\n",
      "alxlnet-base-entity.pb ['Placeholder', 'Placeholder_1', 'Placeholder_2']\n",
      "albert-tiny-entity.pb ['Placeholder', 'Placeholder_1']\n",
      "tiny-bert-entity.pb ['Placeholder', 'Placeholder_1']\n",
      "bert-base-entity.pb ['Placeholder', 'Placeholder_1']\n",
      "albert-base-entity.pb ['Placeholder', 'Placeholder_1']\n"
     ]
    }
   ],
   "source": [
    "transforms = ['add_default_attributes',\n",
    "             'remove_nodes(op=Identity, op=CheckNumerics, op=Dropout)',\n",
    "             'fold_batch_norms',\n",
    "             'fold_old_batch_norms',\n",
    "             'quantize_weights(fallback_min=-10, fallback_max=10)',\n",
    "             'strip_unused_nodes',\n",
    "             'sort_by_execution_order']\n",
    "\n",
    "for pb in pbs:\n",
    "    input_graph_def = tf.GraphDef()\n",
    "    with tf.gfile.FastGFile(pb, 'rb') as f:\n",
    "        input_graph_def.ParseFromString(f.read())\n",
    "    \n",
    "    if 'bert' in pb:\n",
    "        inputs = ['Placeholder', 'Placeholder_1']\n",
    "        outputs = ['dense/BiasAdd']\n",
    "    if 'xlnet'in pb:\n",
    "        inputs = ['Placeholder', 'Placeholder_1', 'Placeholder_2']\n",
    "        outputs = ['transpose_3']\n",
    "        \n",
    "    print(pb, inputs)\n",
    "    \n",
    "    transformed_graph_def = TransformGraph(input_graph_def, \n",
    "                                           inputs,\n",
    "                                           ['logits'] + outputs, transforms)\n",
    "    \n",
    "    with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:\n",
    "        f.write(transformed_graph_def.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_graph(frozen_graph_filename, **kwargs):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "\n",
    "    # https://github.com/onnx/tensorflow-onnx/issues/77#issuecomment-445066091\n",
    "    # to fix import T5\n",
    "    for node in graph_def.node:\n",
    "        if node.op == 'RefSwitch':\n",
    "            node.op = 'Switch'\n",
    "            for index in xrange(len(node.input)):\n",
    "                if 'moving_' in node.input[index]:\n",
    "                    node.input[index] = node.input[index] + '/read'\n",
    "        elif node.op == 'AssignSub':\n",
    "            node.op = 'Sub'\n",
    "            if 'use_locking' in node.attr:\n",
    "                del node.attr['use_locking']\n",
    "        elif node.op == 'AssignAdd':\n",
    "            node.op = 'Add'\n",
    "            if 'use_locking' in node.attr:\n",
    "                del node.attr['use_locking']\n",
    "        elif node.op == 'Assign':\n",
    "            node.op = 'Identity'\n",
    "            if 'use_locking' in node.attr:\n",
    "                del node.attr['use_locking']\n",
    "            if 'validate_shape' in node.attr:\n",
    "                del node.attr['validate_shape']\n",
    "            if len(node.input) == 2:\n",
    "                node.input[0] = node.input[1]\n",
    "                del node.input[1]\n",
    "\n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = load_graph('xlnet-base-entity.pb.quantized')\n",
    "x = g.get_tensor_by_name('import/Placeholder:0')\n",
    "x_len = g.get_tensor_by_name('import/Placeholder_1:0')\n",
    "x_len2 = g.get_tensor_by_name('import/Placeholder_2:0')\n",
    "logits = g.get_tensor_by_name('import/logits:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x, x_len, logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.58 s, sys: 615 ms, total: 3.19 s\n",
      "Wall time: 2.68 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[2, 2, 2, 0, 0]], dtype=int32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "test_sess.run(logits, feed_dict = {x: [[1,2,3,3,4]], x_len: [[1,1,1,1,1]],\n",
    "                                  x_len2: [[1,1,1,1,1]]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# test_sess.run(logits, feed_dict = {x: [[1,2,3,3,4]], x_len: [[1,1,1,1,1]]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['bert-base-entity.pb.quantized',\n",
       " 'tiny-bert-entity.pb.quantized',\n",
       " 'alxlnet-base-entity.pb.quantized',\n",
       " 'xlnet-base-entity.pb.quantized',\n",
       " 'albert-tiny-entity.pb.quantized',\n",
       " 'albert-base-entity.pb.quantized']"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quantized = glob('*.pb.quantized')\n",
    "quantized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm *.pb*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(\n",
    "#     graph_def_file='test.pb',\n",
    "#     input_arrays=['Placeholder', 'Placeholder_1'],\n",
    "#     input_shapes={'Placeholder' : [None, 512], 'Placeholder_1': [None, 512]},\n",
    "#     output_arrays=['logits'],\n",
    "# )\n",
    "# # converter.allow_custom_ops=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]\n",
    "# converter.target_spec.supported_types = [tf.float16]\n",
    "# converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "# converter.experimental_new_converter = True\n",
    "# tflite_model = converter.convert()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, \n",
    "#                                        tf.lite.OpsSet.SELECT_TF_OPS]\n",
    "# converter.target_spec.supported_types = [tf.float16]\n",
    "# converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "# tflite_model = converter.convert()\n",
    "\n",
    "# with open('tiny-bert-sentiment-float16.tflite', 'wb') as f:\n",
    "#     f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, \n",
    "#                                        tf.lite.OpsSet.SELECT_TF_OPS]\n",
    "# converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]\n",
    "# tflite_model = converter.convert()\n",
    "\n",
    "# with open('tiny-bert-sentiment-hybrid.tflite', 'wb') as f:\n",
    "#     f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# interpreter = tf.lite.Interpreter(model_path='tiny-bert-sentiment-hybrid.tflite')\n",
    "# interpreter.allocate_tensors()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
